use std::collections::{BTreeMap, HashMap};

use futures_util::{stream::FuturesUnordered, StreamExt};
use log::info;
use ruma::{
    api::{
        client::keys::{get_keys, upload_keys},
        federation,
    },
    OwnedDeviceId, OwnedUserId, UserId,
};
use serde_json::json;

use crate::{server::server, Error, Result, Ruma};

/// # `POST /_matrix/client/r0/keys/upload`
///
/// Publish end-to-end encryption keys for the sender device.
///
/// - Adds one time keys
/// - If there are no device keys yet: Adds device keys (TODO: merge with
///   existing keys?)
///
pub async fn upload_keys_route(
    body: Ruma<upload_keys::v3::Request>,
) -> Result<upload_keys::v3::Response> {
    let sender_user = body.sender_user.as_ref().expect("user is authenticated");
    let sender_device = body.sender_device.as_ref().expect("user is authenticated");

    // for (key_key , key_value) in &body.one_time_keys {
    //     server()
    //     .users
    //     .get_device_keys(user_id, device_id)?;
    // }

    if let Some(device_keys) = body.device_keys.as_ref() {
        if server()
            .users
            .get_device_keys(&sender_user, &sender_device)?
            .is_none()
        {
            server()
                .users
                .add_device_keys(&sender_user, &sender_device, device_keys)?;
        }
    }
    // Assume that the one time keys do not exist yet
    Ok(upload_keys::v3::Response {
        one_time_key_counts: BTreeMap::new(),
    })
}

/// # `POST /_matrix/client/r0/keys/query`
///
/// Get end-to-end encryption keys for the given users.
///
/// - Always fetches users from other servers over federation
/// - Gets master keys, self-signing keys, user signing keys and device keys.
/// - The master and self-signing keys contain signatures that the user is
///   allowed to see
pub async fn get_device_route(body: Ruma<get_keys::v3::Request>) -> Result<get_keys::v3::Response> {
    let sender_user = body.sender_user.as_ref().expect("user is authenticated");

    get_keys_helper(
        Some(sender_user),
        &body.device_keys,
        |u| u == sender_user,
        true, // Always allow local users to see device names of other local users
    )
    .await
}

pub async fn get_keys_helper<F: Fn(&UserId) -> bool + Send>(
    sender_user: Option<&UserId>,
    device_keys_input: &BTreeMap<OwnedUserId, Vec<OwnedDeviceId>>,
    allowed_signatures: F,
    include_display_names: bool,
) -> Result<get_keys::v3::Response> {
    let mut master_keys = BTreeMap::new();
    let mut self_signing_keys = BTreeMap::new();
    let mut user_signing_keys = BTreeMap::new();
    let mut device_keys = BTreeMap::new();

    let mut get_over_federation = HashMap::new();

    for (user_id, device_ids) in device_keys_input {
        let user_id: &UserId = user_id;

        if !server().config.user_is_local(user_id) {
            get_over_federation
                .entry(user_id.server_name())
                .or_insert_with(Vec::new)
                .push((user_id, device_ids));
            continue;
        }

        if device_ids.is_empty() {
            let mut container = BTreeMap::new();
            for device_id in server().users.all_device_ids(user_id)? {
                if let Some(mut keys) = server().users.get_device_keys(user_id, &device_id)? {
                    // let metadata = server()
                    //     .users
                    //     .get_device_metadata(user_id, &device_id)?
                    //     .ok_or_else(|| {
                    //         Error::bad_database("all_device_keys contained nonexistent device.")
                    //     })?;

                    // add_unsigned_device_display_name(&mut keys, metadata, include_display_names)
                    //     .map_err(|_| Error::bad_database("invalid device keys in database"))?;

                    container.insert(device_id, keys);
                }
            }
            device_keys.insert(user_id.to_owned(), container);
        } else {
            for device_id in device_ids {
                let mut container = BTreeMap::new();
                if let Some(mut keys) = server().users.get_device_keys(user_id, device_id)? {
                    // let metadata = services
                    //     .users
                    //     .get_device_metadata(user_id, device_id)?
                    //     .ok_or(Error::BadRequest(
                    //         ErrorKind::InvalidParam,
                    //         "Tried to get keys for nonexistent device.",
                    //     ))?;

                    // add_unsigned_device_display_name(&mut keys, metadata, include_display_names)
                    //     .map_err(|_| Error::bad_database("invalid device keys in database"))?;
                    container.insert(device_id.to_owned(), keys);
                }
                device_keys.insert(user_id.to_owned(), container);
            }
        }

        info!("Fetching keys for {}", user_id);
        if let Some(master_key) =
            server()
                .users
                .get_master_key(sender_user, user_id, &allowed_signatures)?
        {
            master_keys.insert(user_id.to_owned(), master_key);
        }
        if let Some(self_signing_key) =
            server()
                .users
                .get_self_signing_key(sender_user, user_id, &allowed_signatures)?
        {
            self_signing_keys.insert(user_id.to_owned(), self_signing_key);
        }
        if Some(user_id) == sender_user {
            if let Some(user_signing_key) = server().users.get_user_signing_key(user_id)? {
                user_signing_keys.insert(user_id.to_owned(), user_signing_key);
            }
        }
        info!("Fetched keys for {}", user_id);
    }

    let mut failures = BTreeMap::new();

    // let back_off = |id| async {
    //     match server()
    //         .globals
    //         .bad_query_ratelimiter
    //         .write()
    //         .expect("locked")
    //         .entry(id)
    //     {
    //         hash_map::Entry::Vacant(e) => {
    //             e.insert((Instant::now(), 1));
    //         }
    //         hash_map::Entry::Occupied(mut e) => {
    //             *e.get_mut() = (Instant::now(), e.get().1.saturating_add(1));
    //         }
    //     }
    // };

    info!("Fetching keys over federation");
    let mut futures: FuturesUnordered<_> = get_over_federation
        .into_iter()
        .map(|(server_name, vec)| async move {
            // backoff
            if rand::random() {
                return (
                    server_name,
                    Err(Error::BadServerResponse(
                        "bad query from {server_name:?}, still backing off",
                    )),
                );
            }

            let mut device_keys_input_fed = BTreeMap::new();
            for (user_id, keys) in vec {
                device_keys_input_fed.insert(user_id.to_owned(), keys.clone());
            }

            let request = federation::keys::get_keys::v1::Request {
                device_keys: device_keys_input_fed,
            };
            let response = server()
                .sending
                .send_federation_request(server_name, request)
                .await;

            (server_name, Ok(response))
        })
        .collect();

    while let Some((server_name, response)) = futures.next().await {
        if let Ok(Ok(response)) = response {
            for (user, masterkey) in response.master_keys {
                let (master_key_id, mut master_key) =
                    server().users.parse_master_key(&user, &masterkey)?;

                if let Some(our_master_key) = server().users.get_key(
                    &master_key_id,
                    sender_user,
                    &user,
                    &allowed_signatures,
                )? {
                    let (_, our_master_key) =
                        server().users.parse_master_key(&user, &our_master_key)?;

                    for (user_id, keys) in our_master_key.signatures.iter() {
                        for (key_id, value) in keys {
                            master_key.signatures.insert_signature(
                                user_id.clone(),
                                key_id.clone(),
                                value.clone(),
                            );
                        }
                    }
                }
                let json = serde_json::to_value(master_key).expect("to_value always works");
                let raw = serde_json::from_value(json).expect("Raw::from_value always works");
                server().users.add_cross_signing_keys(
                    &user, &raw, &None, &None,
                    false, /* Dont notify. A notification would trigger another key request resulting in an
                           * endless loop */
                )?;
                master_keys.insert(user.clone(), raw);
            }

            self_signing_keys.extend(response.self_signing_keys);
            device_keys.extend(response.device_keys);
        } else {
            // back_off(server_name.to_owned()).await;
            failures.insert(server_name.to_string(), json!({}));
        }
    }

    Ok(get_keys::v3::Response {
        failures,
        device_keys,
        master_keys,
        self_signing_keys,
        user_signing_keys,
    })
}
